from io import BytesIO

import HighLevelFeatures as HLF
import numpy as np
from PIL import Image


class EvaluateCaloSim:
    def __init__(self, dataset_name, data):
        particle = {
            "1-photons": "photon",
            "1-pions": "pion",
            "2": "electron",
            "3": "electron",
        }[dataset_name]
        self.data = data
        self.hlf = HLF.HighLevelFeatures(
            particle,
            filename="binning_dataset_{}.xml".format(dataset_name.replace("-", "_")),
        )
        self.hlf.CalculateFeatures(data)

    def DrawAverageShower(self):
        buf = BytesIO()
        _ = self.hlf.DrawAverageShower(
            self.data,
            filename=buf,
            title="Shower average",
        )
        buf.seek(0)
        image = Image.open(buf)
        image = np.array(image)
        return image
